import os
import shutil
import os
import shutil
from concurrent.futures import ThreadPoolExecutor
from PIL import Image
import io

base_destination_root = "./"
base_pics_root = "/Volumes/AI-NAS/data/img"


def copy_image(image_info, compress = True, quality=85):
    line_idx, image_path, destination_directory = image_info
    if not image_path:
        return  # 跳过空行

    # 计算目标路径
    destination_path = f"{base_destination_root}{destination_directory}{image_path}"
    # 创建目标路径的文件夹
    os.makedirs(os.path.dirname(destination_path), exist_ok=True)

    source_image_path = f"{base_pics_root}{image_path}"

    try:
        # 复制文件
        shutil.copy(source_image_path, destination_path)

        with open(destination_path, "rb") as f:
            image_data = f.read()

        if compress:
            image_data, original_size, compressed_size = compress_image(image_data, quality)
        with open(destination_path, 'wb') as f:
            f.write(image_data)

        print(f"{line_idx} :  {source_image_path} to {destination_path}")
    except Exception as e:
        print(f"Error copying {source_image_path}: {str(e)}")


def compress_image(image_data, quality=85):
    """
    压缩图片质量
    """
    try:
        img = Image.open(io.BytesIO(image_data))
        if img.mode != 'RGB':
            img = img.convert('RGB')

        original_size = len(image_data)
        output_buffer = io.BytesIO()
        img.save(output_buffer, format='JPEG', quality=quality, optimize=True)
        compressed_data = output_buffer.getvalue()
        compressed_size = len(compressed_data)

        return compressed_data, original_size, compressed_size
    except Exception as e:
        print(f"图片压缩失败: {e}")
        return image_data, len(image_data), len(image_data)  # 如果压缩失败，返回原始数据



def copy_nas_images(source_file, destination_directory, max_workers=15):
    # 确保目标目录存在
    os.makedirs(destination_directory, exist_ok=True)

    # 准备任务列表
    tasks = []
    line_idx = 0

    # 读取所有任务
    with open(source_file, 'r') as file:
        for line in file:
            image_path = line.strip()  # 去除行末的换行符和空格
            line_idx += 1
            tasks.append((line_idx, image_path, destination_directory))

    # 使用线程池执行
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        executor.map(copy_image, tasks)


# 使用示例
source_file = 'pics2.txt'  # 包含图片相对路径的文件
destination_directory = '003_flaws'  # 目标文件夹

copy_nas_images(source_file, destination_directory)